from typing import List, Dict
import json
import os
import time
import requests
from pathlib import Path

from mm_story_agent.prompts_en import role_extract_system, role_review_system, \
    story_to_image_reviser_system, story_to_image_review_system
from mm_story_agent.base import register_tool, init_tool_instance


@register_tool("story_diffusion_t2i")
class StoryDiffusionAgent:
    """
    使用通义万相API生成图像（完全修复版）
    """

    def __init__(self, cfg) -> None:
        self.cfg = cfg
        # 优先从配置读取 API Key，如果没有则从环境变量读取
        self.api_key = cfg.get("api_key") or os.environ.get('DASHSCOPE_API_KEY')
        if not self.api_key:
            raise ValueError("请设置DASHSCOPE_API_KEY环境变量或在配置文件中添加api_key字段")
        # 设置到环境变量，供 QwenAgent 使用
        os.environ["DASHSCOPE_API_KEY"] = self.api_key
        
        self.api_url = "https://dashscope.aliyuncs.com/api/v1/services/aigc/text2image/image-synthesis"
        self.model = cfg.get("model", "wanx-v1")
        
        # 严格的限流配置（适配免费版API）
        self.request_interval = 10.0  # 每次请求间隔10秒
        self.max_retries = 3  # 减少重试次数，避免浪费配额
        self.retry_delay = 30.0  # 限流后等待30秒
        self.task_check_interval = 3.0  # 任务状态检查间隔3秒
        self.max_task_wait = 180  # 单个任务最多等待3分钟
        
        print(f"✅ 图像生成器初始化完成 - 使用API模式: {self.model}")
        print(f"   ⚠️  免费版限流配置:")
        print(f"   - 请求间隔: {self.request_interval}秒（避免限流）")
        print(f"   - 任务轮询间隔: {self.task_check_interval}秒")
        print(f"   - 最大重试: {self.max_retries}次")
        
    def call(self, params: Dict):
        """主调用方法：生成所有页面的图像"""
        pages: List = params["pages"]
        save_path: Path = params["save_path"]
        
        print("\n" + "="*50)
        print("开始图像生成流程")
        print("="*50)
        
        # 1. 提取角色信息
        print("\n📝 步骤1: 提取故事中的角色...")
        role_dict = self.extract_role_from_story(pages)
        print(f"✅ 识别到 {len(role_dict)} 个角色")
        for role, desc in role_dict.items():
            print(f"   - {role}: {desc[:50]}...")
        
        # 2. 生成图像提示词
        print("\n🎨 步骤2: 为每个页面生成图像提示词...")
        image_prompts = self.generate_image_prompt_from_story(pages)
        print(f"✅ 生成了 {len(image_prompts)} 个图像提示词")
        
        # 3. 优化提示词
        print("\n✨ 步骤3: 优化提示词...")
        image_prompts_with_role_desc = []
        for image_prompt in image_prompts:
            for role, role_desc in role_dict.items():
                if role in image_prompt:
                    image_prompt = image_prompt.replace(role, role_desc)
            image_prompts_with_role_desc.append(image_prompt)
        
        # 4. 生成图像（带完善的错误处理）
        print("\n🖼️ 步骤4: 调用API生成图像...")
        print(f"   ⚠️  免费版API限流严格，每张图片间隔 {self.request_interval} 秒")
        print(f"   ⏱️  预计总耗时: {len(image_prompts_with_role_desc) * (self.request_interval + 10)} 秒")
        
        images = self.generate_images_with_api(
            prompts=image_prompts_with_role_desc,
            save_path=save_path,
            style_name=params.get("style_name", "Storybook")
        )
        
        # 统计成功和失败
        success_count = sum(1 for img in images if not img.endswith("_error.png"))
        print("\n" + "="*50)
        print(f"✅ 图像生成完成！")
        print(f"   - 成功: {success_count}/{len(images)} 张")
        print(f"   - 失败: {len(images) - success_count}/{len(images)} 张")
        print("="*50 + "\n")
        
        return {
            "prompts": image_prompts_with_role_desc,
            "generation_results": images,
        }
    
    def generate_images_with_api(self, prompts: List[str], save_path: Path, style_name: str) -> List[str]:
        """批量生成图像（完善的重试和限流控制）"""
        os.makedirs(save_path, exist_ok=True)
        
        images = []
        total = len(prompts)
        
        for i, prompt in enumerate(prompts, 1):
            print(f"\n{'='*60}")
            print(f"[{i}/{total}] 正在生成第 {i} 张图片...")
            print(f"{'='*60}")
            
            retry_count = 0
            success = False
            image_path = None
            
            while retry_count < self.max_retries and not success:
                try:
                    optimized_prompt = self.optimize_prompt(prompt, style_name)
                    print(f"📝 提示词: {optimized_prompt[:100]}...")
                    
                    # 调用API并等待完成
                    image_path = self.call_dashscope_api_with_wait(
                        prompt=optimized_prompt,
                        save_path=save_path,
                        index=i
                    )
                    
                    images.append(image_path)
                    print(f"✅ 成功保存: {image_path}")
                    success = True
                    
                except Exception as e:
                    error_msg = str(e)
                    retry_count += 1
                    
                    print(f"❌ 第 {retry_count} 次尝试失败: {error_msg[:100]}")
                    
                    if "429" in error_msg or "Throttling" in error_msg:
                        if retry_count < self.max_retries:
                            wait_time = self.retry_delay * retry_count
                            print(f"⏳ API限流，等待 {wait_time} 秒后重试...")
                            time.sleep(wait_time)
                        else:
                            print(f"❌ 达到最大重试次数，放弃此图片")
                    else:
                        if retry_count < self.max_retries:
                            print(f"⏳ {self.retry_delay} 秒后重试...")
                            time.sleep(self.retry_delay)
                        else:
                            print(f"❌ 达到最大重试次数，放弃此图片")
            
            # 失败则添加占位符
            if not success:
                error_path = str(save_path / f"p{i:04d}_error.png")
                images.append(error_path)
                print(f"⚠️  使用错误占位符: {error_path}")
            
            # 请求间隔（即使失败也要等待，避免继续限流）
            if i < total:
                print(f"\n⏳ 等待 {self.request_interval} 秒再处理下一张...")
                time.sleep(self.request_interval)
        
        return images
    
    def call_dashscope_api_with_wait(self, prompt: str, save_path: Path, index: int) -> str:
        """
        调用API并等待任务完成（修复异步任务处理）
        """
        # 1. 提交任务
        print("🚀 提交生成任务...")
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json",
            "X-DashScope-Async": "enable"
        }
        
        data = {
            "model": self.model,
            "input": {"prompt": prompt},
            "parameters": {"size": "1024*1024", "n": 1}
        }
        
        response = requests.post(self.api_url, headers=headers, json=data, timeout=30)
        
        if response.status_code != 200:
            raise Exception(f"API请求失败: {response.status_code} - {response.text}")
        
        result = response.json()
        
        if "output" not in result or "task_id" not in result["output"]:
            raise Exception(f"API响应格式错误: {result}")
        
        task_id = result["output"]["task_id"]
        print(f"📋 任务ID: {task_id}")
        
        # 2. 等待任务完成（修复的核心部分）
        image_url = self.wait_for_task(task_id)
        
        # 3. 下载图像
        print("📥 下载图像...")
        image_path = self.download_image(image_url, save_path, index)
        
        return image_path
    
    def wait_for_task(self, task_id: str) -> str:
        """
        等待异步任务完成（修复版 - 正确处理PENDING状态）
        """
        query_url = f"https://dashscope.aliyuncs.com/api/v1/tasks/{task_id}"
        headers = {"Authorization": f"Bearer {self.api_key}"}
        
        start_time = time.time()
        check_count = 0
        last_status = None
        
        print(f"⏳ 等待任务完成（最多 {self.max_task_wait} 秒）...")
        
        while time.time() - start_time < self.max_task_wait:
            check_count += 1
            
            try:
                response = requests.get(query_url, headers=headers, timeout=10)
                
                if response.status_code != 200:
                    print(f"   ⚠️  查询失败: {response.status_code}，继续重试...")
                    time.sleep(self.task_check_interval)
                    continue
                
                result = response.json()
                
                if "output" not in result:
                    print(f"   ⚠️  响应格式异常，继续重试...")
                    time.sleep(self.task_check_interval)
                    continue
                
                task_status = result["output"].get("task_status", "UNKNOWN")
                
                # 显示状态变化
                if task_status != last_status:
                    elapsed = int(time.time() - start_time)
                    print(f"   [{elapsed}s] 状态: {task_status} (第 {check_count} 次检查)")
                    last_status = task_status
                
                # 根据状态处理
                if task_status == "SUCCEEDED":
                    print(f"   ✅ 任务成功完成！")
                    if "results" not in result["output"] or not result["output"]["results"]:
                        raise Exception("任务成功但未返回图像URL")
                    return result["output"]["results"][0]["url"]
                
                elif task_status == "FAILED":
                    error_code = result["output"].get("code", "UNKNOWN")
                    error_msg = result["output"].get("message", "Unknown error")
                    raise Exception(f"任务失败 [{error_code}]: {error_msg}")
                
                elif task_status in ["PENDING", "RUNNING"]:
                    # 关键修复：这些状态是正常的，继续等待
                    time.sleep(self.task_check_interval)
                    continue
                
                else:
                    # 未知状态，继续等待
                    print(f"   ⚠️  未知状态: {task_status}，继续等待...")
                    time.sleep(self.task_check_interval)
                    
            except requests.RequestException as e:
                print(f"   ⚠️  网络错误: {e}，继续重试...")
                time.sleep(self.task_check_interval)
        
        # 超时
        elapsed = int(time.time() - start_time)
        raise Exception(f"任务超时: 等待 {elapsed} 秒仍未完成（检查了 {check_count} 次）")
    
    def download_image(self, image_url: str, save_dir: Path, index: int) -> str:
        """从URL下载图像"""
        try:
            response = requests.get(image_url, timeout=30)
            response.raise_for_status()
            
            image_path = save_dir / f"p{index:04d}.png"
            
            with open(image_path, 'wb') as f:
                f.write(response.content)
            
            # 验证文件大小
            file_size = os.path.getsize(image_path)
            print(f"   💾 文件大小: {file_size / 1024:.1f} KB")
            
            return str(image_path)
            
        except Exception as e:
            raise Exception(f"下载图像失败: {str(e)}")
    
    def optimize_prompt(self, prompt: str, style_name: str) -> str:
        """优化提示词"""
        style_prompts = {
            "Storybook": "cartoon style, storybook illustration, bright colors, children's book style",
            "Japanese Anime": "anime style, Japanese animation, 2D art, beautiful illustration",
            "Digital/Oil Painting": "digital painting, oil painting style, artistic, detailed brushwork",
            "Pixar/Disney Character": "Pixar style, 3D cartoon, Disney animation, vivid characters",
            "Photographic": "photographic style, realistic, high-definition photo, natural lighting",
            "Comic book": "comic book style, American comics, graphic novel, dynamic composition",
            "Line art": "line art style, minimalist lines, black and white illustration",
            "Black and White Film Noir": "black and white film style, film noir, classic cinema",
            "Isometric Rooms": "isometric view, 3D space, miniature scene, game art",
        }
        
        style_desc = style_prompts.get(style_name, "high-quality illustration, detailed")
        return f"{prompt}, {style_desc}, high quality, rich details, best quality"
    
    def extract_role_from_story(self, pages: List) -> Dict:
        """从故事中提取角色"""
        num_turns = self.cfg.get("num_turns", 3)
        
        # 🔥 检测输入文本的语言
        from ..modality_agents.story_agent import detect_language
        sample_text = " ".join(pages[:3]) if len(pages) >= 3 else " ".join(pages)
        language = detect_language(sample_text)
        
        # 根据语言选择 system prompt
        if language == "zh":
            role_extract_system_zh = """从给定的故事内容中提取所有主要角色名称，并生成相应的角色描述。如果有上一轮的结果和改进建议，请根据建议改进之前的角色描述。

## 步骤
1. 首先，识别故事中主要角色的名称。
2. 然后，识别其他频繁出现的角色。
3. 为这些角色生成描述。确保描述**简洁**，重点关注**视觉**特征，如性别或物种，例如"小男孩"或"小鸟"。
4. 确保描述不超过20个字。

## 输入格式
输入包含故事内容，以及可能的上轮输出结果和改进建议，格式如下：
{
    "story_content": "xxx",
    "previous_result": {
        "(角色1的名称)": "xxx",
        "(角色2的名称)": "xxx"
    }, // 空表示第一轮
    "improvement_suggestions": "xxx" // 空表示第一轮
}

## 输出格式
按照以下格式输出角色名称和描述：
{
    "(角色1的名称)": "xxx",
    "(角色2的名称)": "xxx"
}
严格遵循上述步骤，直接输出结果，不要添加任何其他内容。"""
            
            role_review_system_zh = """审查与给定故事对应的角色描述。如果满足要求，输出"Check passed."。如果不满足，提供改进建议。

## 角色描述要求
1. 描述必须**简洁**，**视觉化**，表明性别或物种，例如"小男孩"或"小鸟"。
2. 描述不应包含外观以外的任何信息，如性格或行为。
3. 每个角色的描述不得超过20个字。

## 输入格式
输入包含故事内容和角色提取结果，格式如下：
{
    "story_content": "xxx",
    "role_descriptions": {
        "(角色1的名称)": "xxx",
        "(角色2的名称)": "xxx"
    }
}

## 输出格式
如果不满足要求，直接输出改进建议，不要添加任何其他内容。否则，输出"Check passed."。"""
            
            role_extract_sys = role_extract_system_zh
            role_review_sys = role_review_system_zh
        else:
            role_extract_sys = role_extract_system
            role_review_sys = role_review_system
        
        # 创建 LLM 配置的辅助方法
        def get_llm_cfg(system_prompt, track_history=False):
            cfg = {
                "system_prompt": system_prompt,
                "track_history": track_history
            }
            if self.api_key:
                cfg["api_key"] = self.api_key
            return cfg
        
        role_extractor = init_tool_instance({
            "tool": self.cfg.get("llm", "qwen"),
            "cfg": get_llm_cfg(role_extract_sys, track_history=False)
        })
        
        role_reviewer = init_tool_instance({
            "tool": self.cfg.get("llm", "qwen"),
            "cfg": get_llm_cfg(role_review_sys, track_history=False)
        })
        
        roles = {}
        review = ""
        
        for turn in range(num_turns):
            roles_str, success = role_extractor.call(json.dumps({
                "story_content": pages,
                "previous_result": roles,
                "improvement_suggestions": review,
            }, ensure_ascii=False))
            
            try:
                roles = json.loads(roles_str.strip("```json").strip("```"))
            except:
                roles = {}
            
            review, success = role_reviewer.call(json.dumps({
                "story_content": pages,
                "role_descriptions": roles
            }, ensure_ascii=False))
            
            if review == "Check passed.":
                break
        
        return roles

    def generate_image_prompt_from_story(self, pages: List) -> List[str]:
        """生成图像提示词"""
        num_turns = self.cfg.get("num_turns", 3)
        
        # 🔥 检测输入文本的语言
        from ..modality_agents.story_agent import detect_language
        sample_text = " ".join(pages[:3]) if len(pages) >= 3 else " ".join(pages)
        language = detect_language(sample_text)
        print(f"检测到故事语言: {'中文' if language == 'zh' else '英文'}")
        
        # 根据语言选择 system prompt
        if language == "zh":
            story_to_image_reviser_system_zh = """将给定的故事内容转换为图像描述。如果有上一轮的结果和改进建议，请根据建议改进描述。

## 输入格式
输入包含所有故事页面、当前页面、前一张图片的描述（用于保持场景连续性），以及可能的上轮输出结果和改进建议，格式如下：
{
    "all_pages": ["xxx", "xxx"], // 每个元素是一页故事内容
    "current_page": "xxx",
    "previous_image_description": "xxx", // 前一张图片的描述，用于保持场景连续性（第一张图片时为空）
    "previous_result": "xxx", // 如果为空，表示第一轮
    "improvement_suggestions": "xxx" // 如果为空，表示第一轮
}

## 输出格式
输出一个字符串，描述与当前故事内容对应的图像，不要添加任何其他内容。

## 注意事项
1. 保持简洁。关注主要视觉元素，省略细节。
2. 保留视觉元素。只描述静态场景，避免情节细节。
3. 移除非视觉元素。典型的非视觉元素包括对话、思想和情节。
4. 保留角色名称。
5. 🔥 场景连续性：如果提供了前一张图片的描述，确保当前图片与前一张在场景、背景、色调、风格上保持一定的连续性，使相邻图片能够平滑过渡。例如，如果前一张是室内场景，当前也应该是相似的室内场景；如果前一张是白天，当前也应该是白天（除非故事明确说明时间变化）。"""
            
            story_to_image_review_system_zh = """审查与给定故事内容对应的图像描述。如果满足要求，输出"Check passed."。如果不满足，提供改进建议。

## 图像描述要求
1. 保持简洁。关注主要视觉元素，省略细节。
2. 保留视觉元素。只描述静态场景，避免情节细节。
3. 移除非视觉元素。典型的非视觉元素包括对话、思想和情节。
4. 保留角色名称。

## 输入格式
输入包含所有故事内容、当前故事内容和对应的图像描述，格式如下：
{
    "all_pages": ["xxx", "xxx"],
    "current_page": "xxx",
    "image_description": "xxx"
}

## 输出格式
如果不满足要求，直接输出改进建议，不要添加任何其他内容。否则，输出"Check passed."。"""
            
            image_reviser_system = story_to_image_reviser_system_zh
            image_reviewer_system = story_to_image_review_system_zh
        else:
            image_reviser_system = story_to_image_reviser_system
            image_reviewer_system = story_to_image_review_system
        
        # 创建 LLM 配置的辅助方法
        def get_llm_cfg(system_prompt, track_history=False):
            cfg = {
                "system_prompt": system_prompt,
                "track_history": track_history
            }
            if self.api_key:
                cfg["api_key"] = self.api_key
            return cfg
        
        image_prompt_reviser = init_tool_instance({
            "tool": self.cfg.get("llm", "qwen"),
            "cfg": get_llm_cfg(image_reviser_system, track_history=False)
        })
        
        image_prompt_reviewer = init_tool_instance({
            "tool": self.cfg.get("llm", "qwen"),
            "cfg": get_llm_cfg(image_reviewer_system, track_history=False)
        })
        
        image_prompts = []
        previous_image_description = ""  # 🔥 保存前一张图片的描述
        
        for idx, page in enumerate(pages):
            review = ""
            image_prompt = ""
            
            for turn in range(num_turns):
                # 🔥 传递前一张图片的描述，确保场景连续性
                image_prompt, success = image_prompt_reviser.call(json.dumps({
                    "all_pages": pages,
                    "current_page": page,
                    "previous_image_description": previous_image_description,  # 🔥 传递前一张图片描述
                    "previous_result": image_prompt,
                    "improvement_suggestions": review,
                }, ensure_ascii=False))
                
                # 兼容 DashScope 返回对象/字典/限流错误，统一提取文本
                def _coerce_text(resp) -> str:
                    # 🔥 修复：首先检查 None
                    if resp is None:
                        return "[ERROR] API call returned None"
                    
                    # 已是字符串
                    if isinstance(resp, str):
                        return resp
                    
                    # DashScopeResponse: resp.output.text
                    try:
                        output = getattr(resp, "output", None)
                        if output is not None:
                            text = getattr(output, "text", None)
                            if isinstance(text, str):
                                return text
                    except Exception:
                        pass
                    
                    # dict 结构
                    if isinstance(resp, dict):
                        # 🔥 处理账户欠费（Arrearage）和其他错误
                        code = resp.get("code") or resp.get("error", {}).get("code")
                        message = resp.get("message") or resp.get("error", {}).get("message", "")
                        
                        if code:
                            code_str = str(code)
                            if "Throttling" in code_str:
                                return "[ERROR] Throttling: quota exceeded"
                            elif "Arrearage" in code_str:
                                return f"[ERROR] Arrearage: {message or 'Account payment overdue'}"
                            elif code_str:
                                return f"[ERROR] API Error ({code_str}): {message or 'Unknown error'}"
                        
                        text = (
                            resp.get("output", {}).get("text")
                            or resp.get("text")
                            or message
                        )
                        if isinstance(text, str):
                            return text
                    
                    # 兜底
                    try:
                        return str(resp)
                    except Exception:
                        return "[ERROR] Failed to extract text from response"

                image_prompt = _coerce_text(image_prompt).strip()

                # 🔥 处理 API 错误：直接使用当前页内容作为提示词，避免报错中断
                if image_prompt.startswith("[ERROR]"):
                    print(f"⚠️  API 调用失败，使用页面文本作为提示词: {image_prompt}")
                    image_prompt = page[:500] if page else "A story scene"

                if image_prompt.startswith("Image description:"):
                    image_prompt = image_prompt[len("Image description:"):].strip()
                
                review, success = image_prompt_reviewer.call(json.dumps({
                    "all_pages": pages,
                    "current_page": page,
                    "image_description": image_prompt
                }, ensure_ascii=False))
                
                review = _coerce_text(review).strip()
                
                # 🔥 如果 review 也失败，跳过检查，直接使用当前提示词
                if review.startswith("[ERROR]"):
                    print(f"⚠️  审查 API 调用失败，跳过审查: {review}")
                    break

                if review == "Check passed.":
                    break
            
            image_prompts.append(image_prompt)
            # 🔥 更新前一张图片的描述，用于下一张图片的生成
            previous_image_description = image_prompt
        
        return image_prompts